import { Rect, Vec2 } from "cc";
import { ObjPool } from "./ObjPool";

export class QuadTree<T> {
    points: Map<T, any> = new Map();
    divided: boolean = false;
    northWest: QuadTree<T>;
    northEast: QuadTree<T>;
    southWest: QuadTree<T>;
    southEast: QuadTree<T>;

    constructor(
        public boundary: Rect = new Rect(),
        public capacity = 16
    ) { }

    divide(pool: ObjPool<QuadTree<T>>) {
        let x = this.boundary.x;
        let y = this.boundary.y;
        let w = this.boundary.width / 2;
        let h = this.boundary.height / 2;

        this.northWest = pool && pool.get()?.setBoundary(x, y + h, w, h) || new QuadTree(new Rect(x, y + h, w, h));
        this.northEast = pool && pool.get()?.setBoundary(x + w, y + h, w, h) || new QuadTree(new Rect(x + w, y + h, w, h));
        this.southWest = pool && pool.get()?.setBoundary(x, y, w, h) || new QuadTree(new Rect(x, y, w, h));
        this.southEast = pool && pool.get()?.setBoundary(x + w, y, w, h) || new QuadTree(new Rect(x + w, y, w, h));

        this.divided = true;
    }

    setBoundary(x: number, y: number, w: number, h: number) {
        this.boundary.x = x
        this.boundary.y = y
        this.boundary.width = w
        this.boundary.height = h
        return this
    }

    insert(pos: Vec2, obj: T, pool?: ObjPool<QuadTree<T>>) {
        if (!this.boundary.contains(pos)) {
            return false;
        }

        if (this.points.size < this.capacity) {
            this.points.set(obj, pos);
            return true;
        } else {
            if (!this.divided) {
                this.divide(pool);
            }

            return (
                this.northWest.insert(pos, obj, pool) ||
                this.northEast.insert(pos, obj, pool) ||
                this.southWest.insert(pos, obj, pool) ||
                this.southEast.insert(pos, obj, pool)
            );
        }
    }


    query(range: Rect, found: T[] = []) {
        if (!this.boundary.intersects(range)) {
            return found;
        } else {
            this.points.forEach((p, key) => {
                if (range.contains(p)) {
                    found.push(key);
                }
            });
            if (this.divided) {
                this.northWest.query(range, found);
                this.northEast.query(range, found);
                this.southWest.query(range, found);
                this.southEast.query(range, found);
            }
        }
        return found;
    }
}
